{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed Data Classification with NeMo Curator's `ContentTypeClassifier`\n",
    "\n",
    "This notebook demonstrates the use of NeMo Curator's `ContentTypeClassifier`. The [content type classifier](https://huggingface.co/nvidia/content-type-classifier-deberta) is used to categorize documents into one of 11 distinct speech types based on their content. It helps with data annotation, which is useful in data blending for foundation model training. Please refer to the NemoCurator Content Type Classifier DeBERTa Hugging Face page for more information about the content type classifier, including its output labels, here: https://huggingface.co/nvidia/content-type-classifier-deberta.\n",
    "\n",
    "This tutorial requires at least 1 NVIDIA GPU with:\n",
    "  - Volta™ or higher (compute capability 7.0+)\n",
    "  - CUDA 12.x\n",
    "\n",
    "Before running this notebook, see this [Installation Guide](https://docs.nvidia.com/nemo/curator/latest/admin/installation.html#admin-installation) page for instructions on how to install NeMo Curator. Be sure to use an installation method which includes GPU dependencies.\n",
    "\n",
    "For more information about the classifiers, refer to our [Distributed Data Classification](https://docs.nvidia.com/nemo/curator/latest/curate-text/process-data/quality-assessment/distributed-classifier.html) documentation page."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Silence Curator logs via Loguru\n",
    "import os\n",
    "\n",
    "os.environ[\"LOGURU_LEVEL\"] = \"ERROR\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following imports are required for this tutorial:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "from nemo_curator.core.client import RayClient\n",
    "from nemo_curator.pipeline import Pipeline\n",
    "from nemo_curator.stages.text.classifiers import ContentTypeClassifier\n",
    "from nemo_curator.stages.text.io.reader.jsonl import JsonlReader\n",
    "from nemo_curator.stages.text.io.writer.jsonl import JsonlWriter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run a pipeline in NeMo Curator, we must start a Ray cluster. This can be done manually (see the [Ray documentation](https://docs.ray.io/en/latest/ray-core/starting-ray.html)) or with Curator's `RayClient`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    ray_client = RayClient()\n",
    "    ray_client.start()\n",
    "except Exception as e:\n",
    "    msg = f\"Error initializing Ray client: {e}\"\n",
    "    raise RuntimeError(msg) from e"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialize Read, Classification, and Write Stages\n",
    "\n",
    "Functions in NeMo Curator are called stages. For this tutorial, we will initialize 3 stages: a JSONL file reader, a content type classification stage, and a JSONL file writer.\n",
    "\n",
    "For this tutorial, we will create a sample JSONL file to use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "news_example = \"\"\"\n",
    "Brent awarded for leading collaborative efforts and leading SIA International Relations Committee.\n",
    "\n",
    "Mar 20, 2018\n",
    "\n",
    "The Security Industry Association (SIA) will recognize Richard Brent, CEO, Louroe Electronics with the prestigious 2017 SIA Chairman's Award for his work to support leading the SIA International Relations Committee and supporting key government relations initiatives.\n",
    "\n",
    "With his service on the SIA Board of Directors and as Chair of the SIA International Relations Committee, Brent has forged relationships between SIA and agencies like the U.S. Commercial Service. A longtime advocate for government engagement generally and exports specifically, Brent's efforts resulted in the publication of the SIA Export Assistance Guide last year as a tool to assist SIA member companies exploring export opportunities or expanding their participation in trade.\n",
    "\n",
    "SIA Chairman Denis Hébert will present the SIA Chairman's Award to Brent at The Advance, SIA's annual membership meeting, scheduled to occur on Tuesday, April 10, 2018, at ISC West.\n",
    "\n",
    "\"As the leader of an American manufacturing company, I have seen great business opportunities in foreign sales,\" said Brent. \"Through SIA, I have been pleased to extend my knowledge and experience to other companies that can benefit from exporting. And that is the power of SIA: To bring together distinct companies to share expertise across vertical markets in a collaborative fashion. I'm pleased to contribute, and I thank the Chairman for his recognition.\"\n",
    "\n",
    "\"As a member of the SIA Board of Directors, Richard Brent is consistently engaged on a variety of issues of importance to the security industry, particularly related to export assistance programs that will help SIA members to grow their businesses,\" said Hébert. \"His contributions in all areas of SIA programming have been formidable, but we owe him a particular debt in sharing his experiences in exporting. Thank you for your leadership, Richard.\"\n",
    "\n",
    "Hébert will present SIA award recipients, including the SIA Chairman's Award, SIA Committee Chair of the Year Award and Sandy Jones Volunteer of the Year Award, at The Advance, held during ISC West in Rooms 505/506 of the Sands Expo in Las Vegas, Nevada, on Tuesday, April 10, 10:30-11:30 a.m. Find more info and register at https:/​/​www.securityindustry.org/​advance.\n",
    "\n",
    "The Advance is co-located with ISC West, produced by ISC Security Events. Security professionals can register to attend the ISC West trade show and conference, which runs April 10-13, at http:/​/​www.iscwest.com.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "online_comments_example = \"Hi, great video! I am now a subscriber.\"\n",
    "explanatory_articles_example = \"Photosynthesis is the process by which plants convert sunlight into chemical energy. During this process, plants absorb carbon dioxide from the air and water from the soil, producing glucose for energy and releasing oxygen as a byproduct.\"\n",
    "blogs_example = \"This week I finally tried making sourdough bread from scratch, and let me tell you—it was both harder and more rewarding than I expected. Here's what I learned (and what I'll do differently next time).\"\n",
    "boilerplate_content_example = \"All rights reserved. This website and its contents may not be reproduced, distributed, or transmitted in any form without prior written permission from the publisher.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_file_path = \"./input_data_dir\"\n",
    "\n",
    "# Create sample dataset for the tutorial\n",
    "text = [news_example, online_comments_example, explanatory_articles_example, blogs_example, boilerplate_content_example]\n",
    "df = pd.DataFrame({\"text\": text})\n",
    "\n",
    "try:\n",
    "    os.makedirs(input_file_path, exist_ok=True)\n",
    "    df.to_json(input_file_path + \"/data.jsonl\", orient=\"records\", lines=True)\n",
    "except Exception as e:\n",
    "    msg = f\"Error creating input file: {e}\"\n",
    "    raise RuntimeError(msg) from e"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can define the reader stage with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read existing directory of JSONL files\n",
    "read_stage = JsonlReader(input_file_path, files_per_partition=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The classifier stage is broken down under the hood into a tokenizer stage and a model inference stage. Tokenization is run on the CPU while model inference is run on the GPU. This means that behind the scenes, the `ContentTypeClassifier` stage is actually being broken down into 2 stages (some parameters and details omitted to avoid complexity, please refer to the documentation for more details):\n",
    "\n",
    "```python\n",
    "class TokenizerStage:\n",
    "    self.resources = Resources(cpus=1)\n",
    "    self.model_identifier = \"nvidia/content-type-classifier-deberta\"\n",
    "    self.text_field = \"text\"\n",
    "    self.padding_side = \"right\"\n",
    "    ...\n",
    "class ModelStage:\n",
    "    self.resources = Resources(cpus=1, gpus=1)\n",
    "    self.model_identifier = \"nvidia/content-type-classifier-deberta\"\n",
    "    self.model_inference_batch_size = 256\n",
    "    ...\n",
    "```\n",
    "\n",
    "Optionally, the classifier predictions may be filtered to include only texts with values listed in `filter_by`. If the `filter_by` parameter is set, then a third stage is added:\n",
    "\n",
    "```python\n",
    "def filter_by_category(self, value: str) -> bool:\n",
    "    return value in self.filter_by\n",
    "\n",
    "...\n",
    "\n",
    "if self.filter_by is not None and len(self.filter_by) > 0:\n",
    "    self.stages.append(Filter(filter_fn=self.filter_by_category, filter_field=...))\n",
    "```\n",
    "\n",
    "See the [API Reference](https://docs.nvidia.com/nemo/curator/latest/apidocs/stages/stages.text.classifiers.content_type.html#stages.text.classifiers.content_type.ContentTypeClassifier) for more information about the `ContentTypeClassifier` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the content type classifier\n",
    "classifier_stage = ContentTypeClassifier()\n",
    "\n",
    "# If desired, you may filter your dataset with:\n",
    "# classifier_stage = ContentTypeClassifier(filter_by=[\"News\"])  # noqa: ERA001\n",
    "# See full list of labels here: https://huggingface.co/nvidia/content-type-classifier-deberta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can define a stage for writing the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Write results to a directory\n",
    "output_file_path = \"./content_type_classifier_results\"\n",
    "\n",
    "# Use mode=\"overwrite\" to overwrite the output directory if it already exists\n",
    "# This helps to ensure that the correct output is written\n",
    "write_stage = JsonlWriter(output_file_path, mode=\"overwrite\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialize Pipeline\n",
    "\n",
    "In NeMo Curator, we use pipelines to run distributed data workflows using Ray. Pipelines take care of resource allocation and autoscaling to achieve enhanced performance and minimize GPU idleness.\n",
    "\n",
    "For the distributed data classifiers, we are able to achieve speedups by ensuring that model inference is run in parallel across all available GPUs, while other stages such as I/O, tokenization, and filtering are run across all available CPUs. This is possible because Curator pipelines are composable, which allows each stage in a pipeline to run independently and with its own specified hardware resources."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Pipeline(name='classifier_pipeline', stages=[jsonl_reader(JsonlReader), content_type_classifier_deberta_classifier(ContentTypeClassifier), jsonl_writer(JsonlWriter)])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "classifier_pipeline = Pipeline(name=\"classifier_pipeline\", description=\"Run a classifier pipeline\")\n",
    "\n",
    "# Add stages to the pipeline\n",
    "classifier_pipeline.add_stage(read_stage)\n",
    "classifier_pipeline.add_stage(classifier_stage)\n",
    "classifier_pipeline.add_stage(write_stage)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Composability is also what allows a classifier to sit between pre-processing and post-processing stages. Typical text pre-processing add-ons include text normalization (lowercasing, URL/email removal, Unicode cleanup) and language identification and filtering (to keep only target languages). A full pipeline may look something like:\n",
    "\n",
    "```python\n",
    "pipeline = Pipeline(name=\"full_pipeline\")\n",
    "pipeline.add_stage(read_stage)                # reader (JSONL/S3/etc.)\n",
    "pipeline.add_stage(lang_id_stage)             # optional: language filter\n",
    "pipeline.add_stage(classifier_stage)          # classifier\n",
    "pipeline.add_stage(write_stage)               # writer (JSONL/Parquet)\n",
    "```\n",
    "\n",
    "# Run the  Classifier\n",
    "\n",
    "Let's run the full classifier pipeline:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the pipeline\n",
    "result = classifier_pipeline.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the pipeline ran to completion and the result was written to a JSONL file, we can shut down the Ray cluster with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    ray_client.stop()\n",
    "except Exception as e:  # noqa: BLE001\n",
    "    print(f\"Error stopping Ray client: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inspect the Output\n",
    "\n",
    "The write stage returns a list of written files. We can read the output file as a Pandas DataFrame for inspection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>text</th>\n",
       "      <th>content_pred</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>\\nBrent awarded for leading collaborative effo...</td>\n",
       "      <td>News</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Hi, great video! I am now a subscriber.</td>\n",
       "      <td>Online Comments</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Photosynthesis is the process by which plants ...</td>\n",
       "      <td>Explanatory Articles</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>This week I finally tried making sourdough bre...</td>\n",
       "      <td>Blogs</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>All rights reserved. This website and its cont...</td>\n",
       "      <td>Boilerplate Content</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                text          content_pred\n",
       "0  \\nBrent awarded for leading collaborative effo...                  News\n",
       "1            Hi, great video! I am now a subscriber.       Online Comments\n",
       "2  Photosynthesis is the process by which plants ...  Explanatory Articles\n",
       "3  This week I finally tried making sourdough bre...                 Blogs\n",
       "4  All rights reserved. This website and its cont...   Boilerplate Content"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# For simplicity, we take the first written file from the writer stage\n",
    "# In real pipelines, the writer may return multiple files (shards) or objects\n",
    "result_file = result[0].data[0]\n",
    "\n",
    "result_df = pd.read_json(result_file, lines=True)\n",
    "result_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the predictions were generated as expected."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
